import os
import io
import re
import json
import pandas as pd


class TextHelper:
    '''
        读取文本文件，支持些常用的抽取函数。
        TODO: 改为单例模式 √
    '''

    _INSTANCE = None

    def __new__(cls):
        if cls._INSTANCE is None:
            cls._INSTANCE = super().__new__(cls)
        return cls._INSTANCE

    def read_lines(self, infile):
        with open(infile, "r", encoding="utf-8") as f:
            lines = [line.strip() for line in f.readlines()]
        return lines

    def read_file(self, infile):
        with open(infile, "r", encoding="utf-8") as f:
            text = f.read()
        return text

    def read_jsonl(self, infile):
        with open(infile, "r", encoding="utf-8") as f:
            lines = f.readlines()
            js_data = [json.loads(line.strip().rstrip(",")) for line in lines]
            return js_data

    def read_json(self, infile):
        df = pd.read_json(infile)
        js_data = json.loads(df.to_json())
        res = []
        for idx in range(len(df)):
            sample = {}
            for key in js_data.keys():
                sample[key] = js_data[key][str(idx)]
            res.append(sample)
        return res

    def write_json(self, res, json_outfile):
        with open(json_outfile, 'w', encoding='utf-8') as f:
            json.dump(res, f, ensure_ascii=False, indent=2)

    def write_jsonl(self, res, json_outfile):
        res = [json.dumps(data, ensure_ascii=False) + "\n" for data in res]
        with open(json_outfile, 'w', encoding='utf-8') as f:
            f.writelines(res)

    # def write_file(self, text,infile):
    #     with open(infile, "w", encoding="utf-8") as f:
    #         f.write(text)

    def write_lines(self, res_ls, outfile, sep=None):
        if (sep is not None) and isinstance(sep, (str)):
            res_ls = [res + sep for res in res_ls]
        with open(outfile, "w", encoding="utf-8") as f:
            f.writelines(res_ls)

    def write_lines_append(self, res_ls, outfile, sep=None):
        if (sep is not None) and isinstance(sep, (str)):
            res_ls = [res + sep for res in res_ls]
        with open(outfile, "a", encoding="utf-8") as f:
            f.writelines(res_ls)

    def read_paragraphs(self, infile):
        ''' 抽取以\n\n分隔的段落
        '''
        text = self.read_file(infile)
        paragraphs = text.split('\n\n')
        paragraphs = [paragraph.strip() for paragraph in paragraphs if paragraph.strip() != ""]
        return paragraphs

    # def split_paragraphs_by_prefix(self, text, prefix="query"):
    #     '''
    #     按照prefix来切分段落，用于提取qa对，如：
    #     输入：
    #         query：
    #         answer：
    #         query：
    #         answer：
    #     输出：
    #     out = ["query:..answer...","query:..answer..."]
    #
    #     '''
    #     text = text + f"\n{prefix}"  # 防止最后一个匹配不上
    #     # 使用正向先行断言，防止第偶数个query被消耗，参考：https://www.runoob.com/w3cnote/reg-lookahead-lookbehind.html
    #     pattern = rf"({prefix}:[\s\S]+?)(?={prefix})"  # 正向先行断言，匹配左边（先行），以prefix为后缀（正向）
    #     paragraphs = re.findall(pattern, text, re.DOTALL)  # 列表，finditer是返回迭代器，需要for才能得到具体值
    #     return paragraphs

    def split_paragraphs_by_prefix(self, text, prefix="query"):
        '''
        按照prefix来切分段落，用于提取qa对，如：
        输入：
            query：
            answer：
            query：
            answer：
        输出：
        out = ["query:..answer...","query:..answer..."]

        '''

        def find_all_occurrences(text, prefix):
            positions = []
            start = 0
            while True:
                index = text.find(prefix, start)
                if index == -1:
                    break
                positions.append(index)
                start = index + 1
            return positions

        text = text + f"\n{prefix}"  # 最后一个不用取
        paragraphs = []
        start_positions = find_all_occurrences(text, prefix)
        for start_idx, end_idx in zip(start_positions[:-1], start_positions[1:]):  # 除去最后一个
            paragraph = text[start_idx:end_idx]
            paragraphs.append(paragraph.strip())

        return paragraphs
    def get_content_between_a_b(self, a, b, text):
        ''' 返回text中a和b之间的内容 '''
        match = re.search(f"{a}(.*?){b}", text, re.DOTALL)
        if match:
            return match.group(1).strip().lstrip(":").lstrip("：")
        else:
            return ""

    def get_content_after_b(self, b, text):
        ''' 返回text中b之后的内容 '''
        match = re.search(f"{b}(.*?)$", text, re.DOTALL)
        if match:
            return match.group(1).strip().lstrip(":").lstrip("：").strip()
        else:
            return ""

    def join_with_space(self, text):
        text = " ".join(text.strip().split("\n"))
        text = " ".join(text.strip().split("\t"))
        text = " ".join(text.split())
        return text

    def jsonify_list(self, dict_ls):
        res = []
        for d in dict_ls:
            js_data = json.dumps(d, ensure_ascii=False)
            res.append(js_data + "\n")
        return "".join(res)


txthelper = TextHelper()


## copy from /stanford_alpaca's self-instruct
# https://github.com/tatsu-lab/stanford_alpaca/blob/main/utils.py#L148

def _make_w_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f_dirname = os.path.dirname(f)
        if f_dirname != "":
            os.makedirs(f_dirname, exist_ok=True)
        f = open(f, mode=mode, encoding="utf-8")
    return f


def _make_r_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f = open(f, mode=mode, encoding="utf-8")
    return f


def jdump(obj, f, mode="w", indent=4, default=str):
    """Dump a str or dictionary to a file in json format.

    Args:
        obj: An object to be written.
        f: A string path to the location on disk.
        mode: Mode for opening the file.
        indent: Indent for storing json dictionaries.
        default: A function to handle non-serializable entries; defaults to `str`.
    """
    f = _make_w_io_base(f, mode)
    if isinstance(obj, (dict, list)):
        json.dump(obj, f, indent=indent, default=default, ensure_ascii=False)
    elif isinstance(obj, str):
        f.write(obj)
    else:
        raise ValueError(f"Unexpected type: {type(obj)}")
    f.close()


def jload(f, mode="r"):
    """Load a .json file into a dictionary."""
    f = _make_r_io_base(f, mode)
    jdict = json.load(f)
    f.close()
    return jdict
